# 导入所需库
from PIL import Image, ImageFilter, ImageEnhance
import matplotlib.pyplot as plt
import matplotlib
import os
import cv2
import numpy as np

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置黑体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 设置图片目录路径
image_dir = r"J:\test\C001H001S0001"
result_dir = os.path.join(image_dir, "边缘检测结果")
os.makedirs(result_dir, exist_ok=True)

# 获取目录下所有TIF文件
tif_files = [f for f in os.listdir(image_dir) if f.lower().endswith('.tif')]

# 检查是否有TIF文件
if not tif_files:
    print("目录中没有找到TIF格式图片！")
else:
    # 让用户选择处理几张图片
    print(f"找到 {len(tif_files)} 张TIF图片:")
    for i, f in enumerate(tif_files, 1):
        print(f"{i}. {f}")
    
    try:
        selection = input("请输入要处理的图片编号(多个用逗号分隔，全部处理请输入all): ")
        if selection.lower() != 'all':
            selected_indices = [int(i)-1 for i in selection.split(',')]
            tif_files = [tif_files[i] for i in selected_indices]
    except Exception as e:
        print(f"输入错误: {e}, 将处理全部图片")

    # 处理选中的TIF文件
    for tif_file in tif_files:
        # 构建完整文件路径
        file_path = os.path.join(image_dir, tif_file)
        
        try:
            # 1. 读取原始图片
            original_img = Image.open(file_path)
            
            # 2. 边缘检测处理
            # 转换为OpenCV格式并确保是8位灰度图
            cv_img = np.array(original_img)
            
            # 确保图像是8位格式
            if cv_img.dtype != np.uint8:
                cv_img = cv_img.astype(np.uint8)
            
            # 转换为灰度图
            if len(cv_img.shape) == 3:  # 如果是彩色图
                gray = cv2.cvtColor(cv_img, cv2.COLOR_RGB2GRAY)
            else:
                gray = cv_img
            
            # 高斯模糊降噪
            blurred = cv2.GaussianBlur(gray, (5, 5), 0)
            
            # 自适应阈值处理
            thresh = cv2.adaptiveThreshold(blurred, 255, 
                                         cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                                         cv2.THRESH_BINARY, 11, 2)
            
            # Canny边缘检测
            edges = cv2.Canny(thresh, 100, 200)
            
            # 转换回Pillow格式
            edge_img = Image.fromarray(edges)
            
            # 3. 创建对比图
            plt.figure(figsize=(10, 5))
            
            # 显示原始图片
            plt.subplot(1, 2, 1)
            plt.imshow(original_img)
            plt.title('原始图片')
            plt.axis('off')
            
            # 显示边缘检测结果
            plt.subplot(1, 2, 2)
            plt.imshow(edge_img, cmap='gray')
            plt.title('边缘检测结果')
            plt.axis('off')
            
            # 调整布局并显示
            plt.tight_layout()
            
            # 保存对比图(PNG格式)
            output_name = os.path.splitext(tif_file)[0] + "_边缘检测.png"
            output_path = os.path.join(result_dir, output_name)
            plt.savefig(output_path, dpi=300, format='png', bbox_inches='tight')
            print(f"已保存处理结果到: {output_path}")
            
            # 显示图片
            plt.show()
            plt.close()  # 关闭当前图形
            
        except Exception as e:
            print(f"处理文件 {tif_file} 时出错: {str(e)}")
